import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_absolute_percentage_error
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, SimpleRNN
import matplotlib.pyplot as plt

class LoadPredictor:
    def __init__(self):
        self.lr_model = None
        self.lstm_model = None
        self.rnn_model = None
        
    def create_lstm_model(self, input_shape):
        model = Sequential([
            LSTM(64, input_shape=input_shape),
            Dense(32, activation='relu'),
            Dense(16, activation='relu'),
            Dense(8, activation='relu'),
            Dense(1)
        ])
        model.compile(optimizer='adam', loss='mae')
        return model
    
    def create_rnn_model(self, input_shape):
        model = Sequential([
            SimpleRNN(64, input_shape=input_shape),
            Dense(32, activation='relu'),
            Dense(16, activation='relu'),
            Dense(8, activation='relu'),
            Dense(1)
        ])
        model.compile(optimizer='adam', loss='mae')
        return model
    
    def train_models(self, X_train, y_train):
        # 训练线性回归
        self.lr_model = LinearRegression()
        self.lr_model.fit(X_train, y_train)
        
        # 重塑数据为3D格式(samples, timesteps, features)
        X_train_3d = X_train.reshape((X_train.shape[0], 1, X_train.shape[1]))
        
        # 训练LSTM
        self.lstm_model = self.create_lstm_model((1, X_train.shape[1]))
        self.lstm_model.fit(X_train_3d, y_train, epochs=50, batch_size=32, verbose=0)
        
        # 训练RNN
        self.rnn_model = self.create_rnn_model((1, X_train.shape[1]))
        self.rnn_model.fit(X_train_3d, y_train, epochs=50, batch_size=32, verbose=0)
    
    def plot_predictions(self, y_true, predictions, title):
        """绘制预测结果时间序列图"""
        plt.figure(figsize=(15, 5))
        plt.plot(y_true[:100], label='Actual', alpha=0.7)
        
        for model_name, pred in predictions.items():
            plt.plot(pred[:100], label=f'{model_name}', alpha=0.7)
        
        plt.title(title)
        plt.xlabel('Time Steps')
        plt.ylabel('Load')
        plt.legend()
        plt.grid(True)
        plt.savefig(f'{title.lower().replace(" ", "_")}.png')
        plt.close()

    def evaluate_models(self, X_test, y_test):
        """评估所有模型性能"""
        predictions = {}
        
        # 线性回归预测
        lr_pred = self.lr_model.predict(X_test)
        predictions['Linear Regression'] = lr_pred
        
        # LSTM预测
        X_test_3d = X_test.reshape((X_test.shape[0], 1, X_test.shape[1]))
        lstm_pred = self.lstm_model.predict(X_test_3d)
        predictions['LSTM'] = lstm_pred.reshape(-1)
        
        # RNN预测  
        rnn_pred = self.rnn_model.predict(X_test_3d)
        predictions['RNN'] = rnn_pred.reshape(-1)
        
        # 绘制预测结果
        self.plot_predictions(y_test, predictions, 'Load Prediction Comparison')
        
        # 计算性能指标
        results = {
            'Linear Regression': {
                'MAE': mean_absolute_error(y_test, lr_pred),
                'MAPE': mean_absolute_percentage_error(y_test, lr_pred)
            },
            'LSTM': {
                'MAE': mean_absolute_error(y_test, lstm_pred),
                'MAPE': mean_absolute_percentage_error(y_test, lstm_pred)
            },
            'RNN': {
                'MAE': mean_absolute_error(y_test, rnn_pred),
                'MAPE': mean_absolute_percentage_error(y_test, rnn_pred)
            }
        }
        
        return results